import json
from typing import Callable, Awaitable, Any

from tcp_nio_module.log import error


class MessageRouter:
    """
    消息路由器（抽象接口）
    - 不处理具体业务逻辑
    - 支持注册命令处理函数
    """

    def __init__(self):
        # 命令 -> 异步处理函数
        # handler signature: async def handler(addr, payload) -> dict
        self.handlers: dict[str, Callable[[Any, dict], Awaitable[dict]]] = {}

    def register(self, cmd: str, handler: Callable[[Any, dict], Awaitable[dict]]):
        """
        注册命令处理函数
        """
        self.handlers[cmd] = handler

    async def route(self, message: Any, addr):
        """
        路由消息
        - message 可以是 JSON 字符串或 dict
        - 根据 cmd 找到对应 handler 执行
        - 返回 JSON 字符串响应
        """
        try:
            if isinstance(message, str):
                data = json.loads(message)
            elif isinstance(message, dict):
                data = message
            else:
                raise ValueError(f"不支持的消息类型: {type(message)}")

            cmd = data.get("cmd")
            payload = data.get("payload", {})

            handler = self.handlers.get(cmd)
            if handler:
                result = await handler(addr, payload)
                return json.dumps(result)
            else:
                error(f"[Router] 未知命令: {cmd}")
                return json.dumps({"cmd": "error", "msg": "未知命令"})

        except Exception as e:
            print(f"[Router] 解析失败: {e}, message={message}")
            return json.dumps({"cmd": "error", "msg": "解析错误"})